Find K closest elements

Time: O(LogN+K); Space: O(1); medium

Given a sorted array arr, two integers k and x, find the k closest elements to x in the array.

The result should also be sorted in ascending order.

If there is a tie, the smaller elements are always preferred.

Example 1:

Input: arr = [1,2,3,4,5], k = 4, x = 3

Output: [1,2,3,4]

Example 2:

Input: arr = [1,2,3,4,5], k = 4, x = -1

Output: [1,2,3,4]

Example 3:

Input: arr = [1,2,3], k = 3, x = 2

Output: [2,1,3]

Example 4:

Input: arr = [1,4,6,8], k = 3, x = 3

Output: [4,1,6]

Constraints:

  • 1 <= k <= len(arr)

  • 1 <= len(arr) <= 10^4

  • Absolute value of elements in the array and x will not exceed 104

1. Binary Search [O(LogN+K), O(1)]

[1]:

import bisect

class Solution1(object):
    """
    Time: O(LogN+K)
    Space: O(1)
    """
    def findClosestElements(self, arr, k, x):
        """
        :type arr: List[int]
        :type k: int
        :type x: int
        :rtype: List[int]
        """
        i = bisect.bisect_left(arr, x)
        left, right = i-1, i

        while k:
            if right >= len(arr) or \
               (left >= 0 and abs(arr[left]-x) <= abs(arr[right]-x)):
                left -= 1
            else:
                right += 1
            k -= 1

        return arr[left+1:right]

[4]:

s = Solution1()

arr = [1,2,3,4,5]
k = 4
x = 3
assert s.findClosestElements(arr, k, x) == [1,2,3,4]

arr = [1,2,3,4,5]
k = 4
x = -1
assert s.findClosestElements(arr, k, x) == [1,2,3,4]

arr = [1,2,3]
k = 3
x = 2
s.findClosestElements(arr, k, x) == [1,2,3]

arr = [1,4,6,8]
k = 3
x = 3
assert s.findClosestElements(arr, k, x) == [1,4,6]